library("ggplot2")

make_funnel_plot <- function(..., ylab = "") {
  ggplot(...) +
    geom_point(alpha = 0.55) + 
    scale_size_continuous(range = c(0.5, 2)) +
    theme_classic() +
    theme(
      text = element_text(size = 12.5),
      legend.position = "none",
      plot.title = element_text(hjust = 0.5),
      plot.subtitle = element_text(hjust = 0.5),
      plot.caption = element_text(size = 8)
    ) +
    geom_text(aes(x = 1.45, label="\nt = 1.96", y = 0.8), colour = "black", angle = 45, size = 2.5) +
    scale_x_continuous(limits = c(0, 2), breaks = seq(0, 2, 0.5)) +
    scale_y_continuous(limits = c(0, 0.8), breaks = seq(0, 0.8, 0.2)) +
    geom_abline(
      slope = 1 / qnorm(0.975),
      intercept = 0,
      colour = "red",
      linewidth = 0.25
    ) +
    labs(title = "", x = "Point estimate", y = ylab)
}

save_funnel_plot <- function(path, plot) {
  ggsave(path, plot = plot, height = 4.1, width = 4.1, units = "in")
}

reanalysis_res <- read.csv("data/reanalysis-results.csv")

rd_papers <- read.csv("data/cl-rd-papers-meta.csv")
col_to_merge <- c("estimate_code", "bw_select_strat", "beta_est", "se_est", "weight")
reanalysis_res <- merge(
  reanalysis_res,
  subset(rd_papers, data_avail, col_to_merge),
  by = "estimate_code"
)
rm(rd_papers, col_to_merge)

# Weights used in plot
reanalysis_res$plot_weight <- 0.75 * sqrt(reanalysis_res$weight)


# Figure 4: Funnel plots for replicated studies by method of analysis

# Panel A
fig4a <- make_funnel_plot(
  subset(reanalysis_res, bw_select_strat != "cct mse optimal"),
  aes(
    x = I(abs(beta_est) / y0_sd),
    y = I(se_est / y0_sd),
    size = plot_weight
  ),
  ylab = "Standard error"
)
save_funnel_plot("output/figures/figure4a-funnel-orginal.pdf", fig4a)

# Panel B
fig4b <- make_funnel_plot(
  subset(reanalysis_res, bw_select_strat != "cct mse optimal"),
  aes(
    x = I(abs(cct_conven_coef) / y0_sd),
    y = I(cct_conven_se / y0_sd),
    size = plot_weight
  )
)
save_funnel_plot("output/figures/figure4b-funnel-conven.pdf", fig4b)

# Panel C
fig4c <- make_funnel_plot(
  subset(reanalysis_res, bw_select_strat != "cct mse optimal"),
  aes(
    x = I(abs(cct_robust_coef) / y0_sd),
    y = I(cct_robust_se / y0_sd),
    size = plot_weight
  )
)
save_funnel_plot("output/figures/figure4c-funnel-robust.pdf", fig4c)


# Figure S4: Funnel plots, all studies

# Panel A
figS4a <- make_funnel_plot(
  reanalysis_res,
  aes(
    x = I(abs(beta_est) / y0_sd),
    y = I(se_est / y0_sd),
    size = plot_weight
  ),
  ylab = "Standard error"
)
save_funnel_plot("output/figures/figureS4a-funnel-all-orginal.pdf", figS4a)

# Panel B
figS4b <- make_funnel_plot(
  reanalysis_res,
  aes(
    x = I(abs(cct_conven_coef) / y0_sd),
    y = I(cct_conven_se / y0_sd),
    size = plot_weight
  )
)
save_funnel_plot("output/figures/figureS4b-funnel-all-conven.pdf", figS4b)

# Panel C
figS4c <- make_funnel_plot(
  reanalysis_res,
  aes(
    x = I(abs(cct_robust_coef) / y0_sd),
    y = I(cct_robust_se / y0_sd),
    size = plot_weight
  )
)
save_funnel_plot("output/figures/figureS4c-funnel-all-robust.pdf", figS4c)


# Figure S7 & S8: RDHonest funnel plots

figS7 <- make_funnel_plot(
  subset(reanalysis_res, !is.na(rdhonest_tstat)),
  aes(
    x = I(abs(beta_est) / y0_sd),
    y = I(se_est / y0_sd),
    size = plot_weight
  ),
  ylab = "Standard error"
)
save_funnel_plot("output/figures/figureS7-funnel-honestsub-orginal.pdf", figS7)

figS8 <- make_funnel_plot(
  subset(reanalysis_res, !is.na(rdhonest_tstat)),
  aes(
    x = I(abs(rdhonest_coef) / y0_sd),
    y = I(rdhonest_se / y0_sd),
    size = plot_weight
  )
)
save_funnel_plot("output/figures/figureS8-funnel-honestsub-rdhonest.pdf", figS8)
